package com.test.security.springsecuritydemo10.service;

import com.test.security.springsecuritydemo10.dao.JpaTokenRepository;
import com.test.security.springsecuritydemo10.entity.Token;
import java.util.Optional;
import java.util.UUID;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.security.web.csrf.CsrfToken;
import org.springframework.security.web.csrf.CsrfTokenRepository;
import org.springframework.security.web.csrf.DefaultCsrfToken;

/**
 * Created with IntelliJ IDEA.
 *
 * @author： liuziyang
 * @date： 2023/1/3-22:33
 * @description：
 * @modifiedBy：
 * @version: 1.0
 */
public class CustomCsrfTokenRepository implements CsrfTokenRepository {
  @Autowired private JpaTokenRepository jpaTokenRepository;

  @Override
  public CsrfToken generateToken(HttpServletRequest request) {
    String uuid = UUID.randomUUID().toString();
    return new DefaultCsrfToken("X-CSRF-TOKEN", "_csrf", uuid);
  }

  @Override
  public void saveToken(CsrfToken token, HttpServletRequest request, HttpServletResponse response) {
    String identifier = request.getHeader("X-IDENTIFIER");
    final Optional<Token> existingToken = jpaTokenRepository.findTokenByIdentifier(identifier);
    existingToken.ifPresentOrElse(
        t -> {
          t.setToken(token.getToken());
          jpaTokenRepository.save(t);
        },
        () -> {
          Token tokenNew = new Token();
          tokenNew.setToken(token.getToken());
          tokenNew.setIdentifier(identifier);
          jpaTokenRepository.save(tokenNew);
        });
  }

  @Override
  public CsrfToken loadToken(HttpServletRequest request) {
    String identifier = request.getHeader("X-IDENTIFIER");
    final Optional<Token> existingToken = jpaTokenRepository.findTokenByIdentifier(identifier);
    return existingToken
        .map(token -> new DefaultCsrfToken("X-IDENTIFIER", "_csrf", token.getToken()))
        .orElse(null);
  }
}
